import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import mplfinance as mpf
from statsmodels.tsa.arima.model import ARIMA
from sklearn.preprocessing import MinMaxScaler
from keras.models import Sequential
from keras.layers import LSTM, Dense
import warnings

# 设置中文显示
plt.rcParams["font.sans-serif"] = ["Microsoft YaHei"]
plt.rcParams['axes.unicode_minus'] = False

# 创建mplfinance样式
my_style = mpf.make_mpf_style(
    base_mpf_style='charles',
    rc={'font.family': 'Microsoft YaHei',
        'axes.titlesize': 12,
        'axes.labelsize': 10}
)

warnings.filterwarnings('ignore')


# ======================
# 1. 数据加载与预处理
# ======================

def load_data(filepath):
    """加载股票数据"""
    data = pd.read_csv(filepath)
    data['日期'] = pd.to_datetime(data['日期'])
    data.set_index('日期', inplace=True)

    # 处理交易量和涨跌幅
    if data['交易量'].dtype == object:
        data['交易量'] = data['交易量'].str.replace('M', '').astype(float) * 1e6
    if data['涨跌幅'].dtype == object:
        data['涨跌幅'] = data['涨跌幅'].str.replace('%', '').astype(float) / 100

    # 计算历史波动率
    data['波动率'] = data['收盘'].pct_change().rolling(20).std()
    return data.dropna()


# 加载数据
try:
    data = load_data("D:/MathD/中国平安2318历史数据.csv")
    last_close = data['收盘'][-1]
    avg_volatility = data['波动率'][-20:].mean()
    print(f"数据加载成功！最后收盘价：{last_close:.2f}，近期平均波动率：{avg_volatility * 100:.2f}%")
except Exception as e:
    print(f"数据加载失败: {e}")
    exit()


# ======================
# 2. ARIMA模型预测（增加适度波动）
# ======================

def arima_predict(train_data, steps=29):
    """ARIMA预测"""
    print("开始ARIMA模型训练...")
    model = ARIMA(train_data['收盘'], order=(2, 1, 1))  # 使用更稳健的参数
    model_fit = model.fit()

    # 基础预测
    base_pred = model_fit.forecast(steps=steps).values

    # 添加基于历史波动率的随机波动（限制在±10%）
    volatility = train_data['波动率'][-20:].mean() * 1.2
    random_shock = np.random.normal(0, volatility, steps).clip(-0.1, 0.1)

    forecast = base_pred * (1 + random_shock)
    forecast = np.maximum(forecast, last_close * 0.7)  # 确保不低于现价的70%
    print(f"ARIMA预测完成！使用波动率：{volatility * 100:.2f}%")
    return np.round(forecast, 2)


arima_pred = arima_predict(data)


# ======================
# 3. LSTM模型预测（增加适度波动）
# ======================

def lstm_predict(train_data, lookback=60, steps=29):
    """LSTM预测"""
    print("开始LSTM模型训练...")
    scaler = MinMaxScaler()
    scaled_data = scaler.fit_transform(train_data[['收盘']])

    # 准备训练数据
    X, y = [], []
    for i in range(lookback, len(scaled_data)):
        X.append(scaled_data[i - lookback:i, 0])
        y.append(scaled_data[i, 0])

    X, y = np.array(X), np.array(y)
    X = X.reshape(X.shape[0], X.shape[1], 1)

    # 构建模型（使用exponential激活确保正输出）
    model = Sequential()
    model.add(LSTM(64, return_sequences=True, input_shape=(X.shape[1], 1)))
    model.add(LSTM(32))
    model.add(Dense(1, activation='exponential'))  # 关键修改
    model.compile(optimizer='adam', loss='mse')
    model.fit(X, y, epochs=40, batch_size=32, verbose=0)

    # 预测未来（添加趋势波动）
    last_window = scaled_data[-lookback:]
    predictions = []
    current_vol = train_data['波动率'][-1]  # 使用最新波动率

    for i in range(steps):
        x = last_window.reshape(1, lookback, 1)
        pred = model.predict(x, verbose=0)[0, 0]

        # 添加随机波动（限制在±8%）
        pred = pred * (1 + np.random.uniform(-0.08, 0.08))
        predictions.append(pred)
        last_window = np.append(last_window[1:], pred)

    # 修复的逆变换代码（括号闭合问题已解决）
    predictions = scaler.inverse_transform(np.array(predictions).reshape(-1, 1))
    predictions = np.maximum(predictions, last_close * 0.7)  # 确保不低于现价的70%
    print("LSTM预测完成！")
    return np.round(predictions, 2)


lstm_pred = lstm_predict(data)


# ======================
# 4. 生成预测结果（确保价格合理性）
# ======================

def generate_ohlc(pred, last_price):
    """生成OHLC数据（确保价格有效性）"""
    opens = np.zeros_like(pred)
    highs = np.zeros_like(pred)
    lows = np.zeros_like(pred)
    closes = np.round(pred.flatten(), 2)

    opens[0] = last_price
    for i in range(1, len(pred)):
        opens[i] = closes[i - 1]  # 今日开盘=昨日收盘

    # 生成合理波动（3-6%）
    volatility = np.random.uniform(0.03, 0.06, len(pred))

    # 计算价格区间（确保最高>max(开,收)，最低<min(开,收)且>0）
    price_range = np.abs(opens - closes)
    highs = np.maximum(opens, closes) * (1 + volatility * 0.7)
    lows = np.maximum(  # 双重保护确保不为负
        np.minimum(opens, closes) * (1 - volatility * 0.7),
        last_close * 0.6  # 最低不低于现价的60%
    )

    return np.round(opens, 2), np.round(highs, 2), np.round(lows, 2), closes


# 生成预测日期
future_dates = pd.date_range(start='2025-08-08', periods=29)
date_labels = [date.strftime('%Y/%m/%d') for date in future_dates]

# 融合两种预测结果
combined_pred = (lstm_pred.flatten() * 0.6 + arima_pred * 0.4)

# 生成OHLC数据
opens, highs, lows, closes = generate_ohlc(combined_pred, last_close)

# 创建预测DataFrame
predictions = pd.DataFrame({
    'Date': future_dates,
    '日期': date_labels,
    '开盘': opens,
    '最高': highs,
    '最低': lows,
    '收盘': closes
})

# 最终校验
assert all(predictions[['开盘', '最高', '最低', '收盘']].min() > 0), "存在负价格！"

# 保存预测结果
predictions[['日期', '开盘', '最高', '最低', '收盘']].to_csv(
    '中国平安股票预测结果.csv',
    index=False,
    encoding='utf-8-sig'
)

# ======================
# 5. 可视化
# ======================

# K线图数据准备
ohlc_data = predictions.set_index('Date')[['开盘', '最高', '最低', '收盘']]
ohlc_data.columns = ['Open', 'High', 'Low', 'Close']
ohlc_data.index.name = 'Date'

# 创建K线图
print("正在生成K线图...")
mpf.plot(ohlc_data,
         type='candle',
         style=my_style,
         title='中国平安股票价格预测\n2025/08/08-2025/09/05',
         ylabel='价格（元）',
         volume=False,
         savefig='中国平安预测K线图.png',
         figratio=(12, 6),
         figscale=1.2)

# 创建对比图
print("正在生成对比图...")
plt.figure(figsize=(12, 6))
plt.plot(future_dates, lstm_pred, label='LSTM预测', alpha=0.7)
plt.plot(future_dates, arima_pred, label='ARIMA预测', alpha=0.7)
plt.plot(future_dates, combined_pred, label='融合预测', linewidth=2)

# 添加波动区间
plt.fill_between(future_dates,
                 predictions['最低'],
                 predictions['最高'],
                 color='gray', alpha=0.2, label='价格区间')

plt.title('预测结果对比（LSTM 60% + ARIMA 40%）', fontsize=14)
plt.xlabel('日期', fontsize=12)
plt.ylabel('收盘价（元）', fontsize=12)
plt.xticks(rotation=45)
plt.legend()
plt.grid(linestyle='--', alpha=0.6)
plt.tight_layout()
plt.savefig('预测结果对比.png', dpi=300)
plt.close()

print("\n预测完成！结果已保存到：")
print("1. 中国平安股票预测结果.csv")
print("2. 中国平安预测K线图.png")
print("3. 预测结果对比.png")